﻿using System;
using System.Buffers;
using System.Collections.Generic;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using JpegLibrary;
using PooledGrowableBufferHelper;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using SixLabors.ImageSharp.Processing;
using TiffLibrary;
using WholeSlideImageLibrary;
using WholeSlideImageLibrary.DeepZoom;
using WholeSlideImageLibrary.Formats.OptimizedPyramidTiledImage;

namespace WholeSlideImageViewer.OptiExport
{
    internal class OptiExporter : IAsyncDisposable
    {
        // Options
        private readonly int _tileSize;
        private readonly bool _enableTiffCompatibility;

        // Source
        private WholeSlideImageShim? _image;
        private DeepZoomTileSource? _deepZoom;

        // Output
        private readonly string _outputFile;
        private Stream? _stream;

        // Progress report
        private readonly OptiExportProgress _progress;

        public OptiExportProgress Progress => _progress;

        public OptiExporter(IOptiOutputOptions options, string outputFile)
        {
            _tileSize = options.TileSize;
            _enableTiffCompatibility = options.EnableTiffCompatibility;
            _outputFile = outputFile;
            _progress = new OptiExportProgress();
        }

        public async Task InitializeAsync(WholeSlideImage image, CancellationToken cancellationToken)
        {
            _image = WholeSlideImageShim.Wrap(image);
            _deepZoom = await _image.OpenDeepZoomAsync(new DeepZoomProfile((short)_tileSize, 0, EncodedImageFileFormat.Jpeg), cancellationToken);
            _stream = new FileStream(_outputFile, FileMode.Create, FileAccess.ReadWrite, FileShare.None, 4096, true);
        }

        public async Task RunAsync(CancellationToken cancellationToken)
        {
            if (_image is null || _deepZoom is null || _stream is null)
            {
                throw new InvalidOperationException();
            }

            int tileSize = _tileSize;
            var writer = new OptimizedPyramidTiledImageWriter(_stream, true);

            try
            {
                // 文件头部
                await writer.WriteFileHeaderAsync(OptimizedPyramidTiledImageWriter.FileHeaderLength, cancellationToken);
                cancellationToken.ThrowIfCancellationRequested();

                // 跳过 MainBlock 区域
                long mainBlockPointer = await writer.SkipMainBlockAsync(cancellationToken);

                // 获取层数信息，计算Tile数量，并写入Tile Collection区域
                DeepZoomLayer[] dzLayers = DeepZoomHelper.CalculateDeepZoomLayers(_image.Width, _image.Height);

                // 计算Level Collection跳过的长度
                long tileCollectionPointer = writer.BytesWritten + OptimizedPyramidTiledImageWriter.GetLevelCollectionBlockLength(dzLayers.Length);

                // 计算Tile Collection偏移量
                long tileCollectionOffset = tileCollectionPointer;
                long[] tileLevelPointers = new long[dzLayers.Length];
                int[] tileCountPerLevel = new int[dzLayers.Length];
                for (int layer = dzLayers.Length - 1; layer >= 0; layer--)
                {
                    int level = dzLayers.Length - 1 - layer;
                    int tileAcross = (dzLayers[layer].Width + tileSize - 1) / tileSize;
                    int tileDown = (dzLayers[layer].Height + tileSize - 1) / tileSize;
                    tileLevelPointers[level] = tileCollectionOffset;
                    tileCountPerLevel[level] = tileAcross * tileDown;
                    tileCollectionOffset += OptimizedPyramidTiledImageWriter.GetTileCollectionBlockLength(tileCountPerLevel[level]);
                }

                cancellationToken.ThrowIfCancellationRequested();

                // 写入Level Collection
                long levelCollectionPointer = await writer.WriteLevelCollectionAsync(tileLevelPointers, cancellationToken);

                // 跳过Tile Collection区域
                for (int level = 0; level < tileLevelPointers.Length; level++)
                {
                    await writer.SkipTileCollectionAsync(tileCountPerLevel[level], cancellationToken);
                }

                cancellationToken.ThrowIfCancellationRequested();

                // Metadata
                IEnumerable<string>? metadataNames = await _image.GetAllTextPropertiesAsync(cancellationToken);
                var metadataList = new List<FileMetadataItem>();
                foreach (string? name in metadataNames)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    metadataList.Add(await writer.WriteMetadataItemAsync(name, await _image.ReadTextPropertyAsync(name, cancellationToken), cancellationToken));
                }
                long metadataBlockPointer = await writer.WriteMetadataBlockAsync(metadataList.ToArray(), cancellationToken);

                cancellationToken.ThrowIfCancellationRequested();

                // Associated Images
                IEnumerable<string>? associatedImageNames = await _image.GetAllAssociatedImagesAsync(cancellationToken);
                var associatedImages = new List<FileAssociatedImageItem>();
                foreach (string? name in associatedImageNames)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    System.Drawing.Size dimensions = await _image.GetAssociatedImageSizeAsync(name, cancellationToken);
                    using (PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream())
                    {
                        await _image.EncodeAssociatedImageAsync(name, EncodedImageFileFormat.Jpeg, (IBufferWriter<byte>)ms, cancellationToken);
                        using PooledMemoryStream ms2 = Optimize(ms.ToReadOnlySequence());
                        ms2.Seek(0, SeekOrigin.Begin);
                        associatedImages.Add(await writer.WriteAssociatedImageItemAsync(ms, name, dimensions.Width, dimensions.Height, "jpeg", cancellationToken));
                    }
                }
                long associatedImagesBlock = await writer.WriteAssociatedImagesBlockAsync(associatedImages.ToArray(), cancellationToken);

                cancellationToken.ThrowIfCancellationRequested();

                // 写入Tiles之前先计算一下Tile数量
                int tileCount = 0;
                foreach (int levelCount in tileCountPerLevel)
                {
                    tileCount += levelCount;
                }
                _progress.NotifyTileCountEstimated(tileCount);

                // 写入Tiles
                FileContentRegion[][] levels;
                levels = await WriteTilesUsingLevel0Async(writer, _deepZoom, dzLayers, tileSize, cancellationToken);
                await _stream.FlushAsync(cancellationToken);

                cancellationToken.ThrowIfCancellationRequested();

                // 返回写入MainBlock
                writer.Seek(mainBlockPointer);
                await writer.WriteMainBlockAsync(new FileMainBlock(_image.Width, _image.Height, (byte)levels.Length, (ushort)tileSize, "jpeg", 0, 0, levelCollectionPointer, metadataBlockPointer, associatedImagesBlock), cancellationToken);

                cancellationToken.ThrowIfCancellationRequested();

                // 返回写入Tile Collection区域
                writer.Seek(tileCollectionPointer);
                for (int level = 0; level < levels.Length; level++)
                {
                    await writer.WriteTileCollectionAsync(levels[level], cancellationToken);
                }

                await _stream.FlushAsync(cancellationToken);

                if (_enableTiffCompatibility)
                {
                    _stream.Seek(0, SeekOrigin.Begin);
                    await WriteTiffAsync(_stream, tileSize, dzLayers, levels, cancellationToken);
                }

                await _stream.DisposeAsync().ConfigureAwait(false);
                _stream = null;
            }
            catch (OperationCanceledException)
            {
                _progress.NotifyCanceled();
                return;
            }
            catch (Exception)
            {
                _progress.NotifyFailed();
                return;
            }

            _progress.NotifyCompleted();
        }

        private async Task<FileContentRegion[][]> WriteTilesUsingLevel0Async(OptimizedPyramidTiledImageWriter writer, DeepZoomTileSource dz, DeepZoomLayer[] dzLayers, int tileSize, CancellationToken cancellationToken)
        {
            var levels = new FileContentRegion[dzLayers.Length][];

            // 最底层
            int layer = dzLayers.Length - 1;
            int tileAcross = (dzLayers[layer].Width + tileSize - 1) / tileSize;
            int tileDown = (dzLayers[layer].Height + tileSize - 1) / tileSize;
            int level = dzLayers.Length - 1 - layer;

            FileContentRegion[] baseLevelTiles = levels[level] = new FileContentRegion[tileAcross * tileDown];

            int index = 0;
            for (int row = 0; row < tileDown; row++)
            {
                for (int col = 0; col < tileAcross; col++)
                {
                    cancellationToken.ThrowIfCancellationRequested();
                    using (PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream())
                    {
                        await dz.WriteToBufferAsync(layer, col, row, ms, cancellationToken).ConfigureAwait(false);
                        using PooledMemoryStream ms2 = Optimize(ms.ToReadOnlySequence());
                        baseLevelTiles[index++] = await writer.WriteAsync(ms2.ToReadOnlySequence(), cancellationToken).ConfigureAwait(false);
                    }
                    _progress.NotifyOneTileWritten();
                }
            }

            // 逐步往上
            for (layer = dzLayers.Length - 2; layer >= 0; layer--)
            {
                tileAcross = (dzLayers[layer].Width + tileSize - 1) / tileSize;
                tileDown = (dzLayers[layer].Height + tileSize - 1) / tileSize;
                level = dzLayers.Length - 1 - layer;

                FileContentRegion[] currentLevelTiles = levels[level] = new FileContentRegion[tileAcross * tileDown];

                index = 0;
                for (int row = 0; row < tileDown; row++)
                {
                    for (int col = 0; col < tileAcross; col++)
                    {
                        cancellationToken.ThrowIfCancellationRequested();

                        // 计算Tile大小
                        CalculateTileSize(dzLayers[layer], tileSize, row, col, out int tileWidth, out int tileHeight);

                        // 获取下一层对应的区域
                        CalculateTileSize(dzLayers[layer + 1], tileSize * 2, row, col, out int baseWidth, out int baseHeight);
                        int tileCountX = (baseWidth + tileSize - 1) / tileSize;
                        int tileCountY = (baseHeight + tileSize - 1) / tileSize;

                        // 创建Tile图像
                        using (var canvas = new Image<Rgba32>(baseWidth, baseHeight))
                        {
                            // 读取下一层对应区域图像并绘制到canvas上
                            for (int tileX = 0; tileX < tileCountX; tileX++)
                            {
                                for (int tileY = 0; tileY < tileCountY; tileY++)
                                {
                                    cancellationToken.ThrowIfCancellationRequested();
                                    int baseTileAcross = (dzLayers[layer + 1].Width + tileSize - 1) / tileSize;
                                    int tileIndex = (col * 2 + tileX) + baseTileAcross * (row * 2 + tileY);

                                    using (PooledMemoryStream ms = PooledMemoryStreamManager.Shared.GetStream())
                                    {
                                        writer.Seek(baseLevelTiles[tileIndex].Offset);
                                        await writer.ReadToAsync(ms, baseLevelTiles[tileIndex].Length, cancellationToken).ConfigureAwait(false);
                                        ms.Seek(0, SeekOrigin.Begin);
                                        using (Image<Rgba32> baseTile = Image.Load<Rgba32>(ms))
                                        {
                                            canvas.Mutate(ctx => ctx.DrawImage(baseTile, new Point(tileX * tileSize, tileY * tileSize), opacity: 1));
                                        }
                                    }
                                }
                            }

                            // 缩放
                            canvas.Mutate(ctx => ctx.Resize(tileWidth, tileHeight));

                            // 保存到流
                            writer.SeekToEnd();
                            currentLevelTiles[index++] = await OptimizeAndSaveAsync(canvas, writer, cancellationToken).ConfigureAwait(false);
                        }

                        _progress.NotifyOneTileWritten();
                    }
                }

                baseLevelTiles = currentLevelTiles;
            }

            return levels;
        }

        private static void CalculateTileSize(DeepZoomLayer layer, int tileSize, int row, int col, out int tileWidth, out int tileHeight)
        {
            int offsetX = col * tileSize;
            int offsetY = row * tileSize;
            tileWidth = Math.Min(tileSize, layer.Width - offsetX);
            tileHeight = Math.Min(tileSize, layer.Height - offsetY);
        }

        private static ushort[] s_yCbCrSubSampling = new ushort[] { 2, 2 };
        private static ushort[] s_bitsPerSample = new ushort[] { 8, 8, 8 };
        private static TiffRational[] s_defaultReferenceBlackWhite = new TiffRational[]
        {
            new TiffRational(0, 1), new TiffRational(255, 1),
            new TiffRational(128, 1), new TiffRational(255, 1),
            new TiffRational(128, 1), new TiffRational(255, 1)
        };
        private static async Task WriteTiffAsync(Stream stream, int tileSize, DeepZoomLayer[] dzLayers, FileContentRegion[][] levels, CancellationToken cancellationToken)
        {
            using (TiffFileWriter writer = await TiffFileWriter.OpenAsync(stream, leaveOpen: true, useBigTiff: true, cancellationToken: cancellationToken))
            {
                // 移动到文件末尾
                writer.Seek(stream.Length);

                TiffStreamOffset ifdOffset = default;

                // 从最大的层开始
                for (int i = 0; i < dzLayers.Length; i++)
                {
                    FileContentRegion[] level = levels[i];
                    DeepZoomLayer layer = dzLayers[dzLayers.Length - 1 - i];
                    int tileAcross = (layer.Width + tileSize - 1) / tileSize;
                    int tileDown = (layer.Height + tileSize - 1) / tileSize;

                    ulong[] tileOffsets = new ulong[tileAcross * tileDown];
                    ulong[] tileByteCounts = new ulong[tileAcross * tileDown];

                    // 输出Tiles
                    for (int row = 0; row < tileDown; row++)
                    {
                        for (int col = 0; col < tileAcross; col++)
                        {
                            int offset = row * tileAcross + col;
                            FileContentRegion region = level[offset];

                            // 对于图像大小不足tileSize的Tile，生成一张新的Tile
                            int width = Math.Min(tileSize, layer.Width - col * tileSize);
                            int height = Math.Min(tileSize, layer.Height - row * tileSize);
                            if (width < tileSize || height < tileSize)
                            {
                                // 记录当前流的位置
                                long position = stream.Position;

                                // 读取当前图片
                                stream.Seek(region.Offset, SeekOrigin.Begin);
                                byte[] buffer = ArrayPool<byte>.Shared.Rent(region.Length);
                                try
                                {
                                    await stream.ReadAsync(buffer, 0, region.Length, cancellationToken);

                                    using (var image = Image.Load<Rgb24>(new MemoryStream(buffer, 0, region.Length)))
                                    using (var tile = new Image<Rgb24>(tileSize, tileSize))
                                    {
                                        // 绘制到tileSize x tileSize的图片上
                                        tile.Mutate(ctx => ctx.DrawImage(image, opacity: 1));
                                        // 编码
                                        byte[] jpeg = OptimizeToBuffer(tile);
                                        // 恢复流的位置
                                        stream.Seek(position, SeekOrigin.Begin);
                                        // 输出到流中
                                        TiffStreamOffset tiffOffset = await writer.WriteAlignedBytesAsync(jpeg, 0, jpeg.Length, cancellationToken);
                                        level[offset] = region = new FileContentRegion(tiffOffset.Offset, jpeg.Length);
                                    }
                                }
                                finally
                                {
                                    ArrayPool<byte>.Shared.Return(buffer);
                                }
                            }

                            // 记录
                            tileOffsets[offset] = (ulong)region.Offset;
                            tileByteCounts[offset] = (ulong)region.Length;
                        }
                    }

                    // 输出IFD
                    using (TiffImageFileDirectoryWriter ifdWriter = writer.CreateImageFileDirectory())
                    {
                        await ifdWriter.WriteTagAsync(TiffTag.PhotometricInterpretation, new TiffValueCollection<ushort>((ushort)TiffPhotometricInterpretation.YCbCr), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.ImageWidth, new TiffValueCollection<uint>((uint)layer.Width), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.ImageLength, new TiffValueCollection<uint>((uint)layer.Height), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.TileWidth, new TiffValueCollection<uint>((uint)tileSize), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.TileLength, new TiffValueCollection<uint>((uint)tileSize), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.Compression, new TiffValueCollection<ushort>((ushort)TiffCompression.Jpeg), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.SamplesPerPixel, new TiffValueCollection<ushort>(3), cancellationToken);

                        await ifdWriter.WriteTagAsync(TiffTag.BitsPerSample, new TiffValueCollection<ushort>(s_bitsPerSample), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.TileOffsets, new TiffValueCollection<ulong>(tileOffsets), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.TileByteCounts, new TiffValueCollection<ulong>(tileByteCounts), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.YCbCrSubSampling, new TiffValueCollection<ushort>(s_yCbCrSubSampling), cancellationToken);
                        await ifdWriter.WriteTagAsync(TiffTag.ReferenceBlackWhite, new TiffValueCollection<TiffRational>(s_defaultReferenceBlackWhite), cancellationToken);

                        await ifdWriter.WriteTagAsync(TiffTag.NewSubfileType, new TiffValueCollection<uint>(ifdOffset.IsZero ? (uint)TiffNewSubfileType.None : (uint)TiffNewSubfileType.ReducedResolution), cancellationToken);

                        ifdOffset = await ifdWriter.FlushAsync(ifdOffset, cancellationToken);
                    }
                }

                await writer.FlushAsync(cancellationToken);
            }

        }

        #region Optimize

        private static PooledMemoryStream Optimize(ReadOnlySequence<byte> sequence)
        {
            PooledMemoryStream? stream = null;
            try
            {
                var optimizer = new JpegOptimizer();
                optimizer.MostOptimalCoding = true;
                optimizer.SetInput(sequence);
                optimizer.Scan();

                stream = PooledMemoryStreamManager.Shared.GetStream();
                optimizer.SetOutput(stream);
                optimizer.Optimize();

                PooledMemoryStream s = stream;
                stream = null;
                return s;
            }
            finally
            {
                stream?.Dispose();
            }

        }

        private static async Task<FileContentRegion> OptimizeAndSaveAsync(Image canvas, OptimizedPyramidTiledImageWriter writer, CancellationToken cancellationToken)
        {
            cancellationToken.ThrowIfCancellationRequested();

            PooledMemoryStream? ms = PooledMemoryStreamManager.Shared.GetStream();
            PooledMemoryStream? ms2 = null;
            try
            {
                canvas.SaveAsJpeg(ms);
                var optimizer = new JpegOptimizer(16384);
                optimizer.MostOptimalCoding = true;
                optimizer.SetInput(ms.ToReadOnlySequence());
                optimizer.Scan();

                ms2 = PooledMemoryStreamManager.Shared.GetStream();
                optimizer.SetOutput(ms2);
                optimizer.Optimize(true);

                ms.Dispose();
                ms = null;

                return await writer.WriteAsync(ms2.ToReadOnlySequence(), cancellationToken).ConfigureAwait(false);
            }
            finally
            {
                ms?.Dispose();
                ms2?.Dispose();
            }
        }

        private static byte[] OptimizeToBuffer(Image canvas)
        {
            PooledMemoryStream? ms = PooledMemoryStreamManager.Shared.GetStream();
            PooledMemoryStream? ms2 = null;
            try
            {
                canvas.SaveAsJpeg(ms);
                var optimizer = new JpegOptimizer(16384);
                optimizer.MostOptimalCoding = true;
                optimizer.SetInput(ms.ToReadOnlySequence());
                optimizer.Scan();

                ms2 = PooledMemoryStreamManager.Shared.GetStream();
                optimizer.SetOutput(ms2);
                optimizer.Optimize(true);

                ms.Dispose();
                ms = null;

                return ms2.ToArray();
            }
            finally
            {
                ms?.Dispose();
                ms2?.Dispose();
            }
        }

        public async ValueTask DisposeAsync()
        {
            if (!(_stream is null))
            {
                await _stream.DisposeAsync().ConfigureAwait(false);
                File.Delete(_outputFile);
                _stream = null;
            }
        }

        #endregion
    }
}
